#!/bin/bash
# Source this file before running test to setup
#
#   source ./common.env
set -eou pipefail

SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
# Mark all repos as safe in the test context, since wandb fetchs metadata about the repo and it's a
# catch-22 to get the project root and mark it safe if you don't know the project root
PROJECT_ROOT=$(realpath $SCRIPT_DIR/../../../..)

exit_if_max_steps_reached() {
  # Early stopping to save compute if max step has been reached
  STEPS_SO_FAR=$(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS || echo 0)
  if [[ $STEPS_SO_FAR -ge $MAX_STEPS ]]; then
      echo "[INFO] Target step $MAX_STEPS reached, skipping run"
      exit 0
  fi
  echo "[INFO] Steps so far: $STEPS_SO_FAR, running till $MAX_STEPS steps"
}

EXP_NAME=$(basename $0 .sh)
EXP_DIR=$SCRIPT_DIR/$EXP_NAME
LOG_DIR=$EXP_DIR/logs
CKPT_DIR=$EXP_DIR/ckpts
JSON_METRICS=$EXP_DIR/metrics.json
RUN_LOG=$EXP_DIR/run.log

# Test script has path:  tests/test_suites/llm/${EXP_NAME}.sh
# where config has path:  examples/configs/recipes/llm/${EXP_NAME}.yaml
# We will assume/check the path matches this pattern
CONFIG_PATH=$(echo $SCRIPT_DIR/${EXP_NAME}.yaml | sed 's#tests/test_suites#examples/configs/recipes#')
if [[ ! -f $CONFIG_PATH ]]; then
  echo "[ERROR] Config file $CONFIG_PATH not found"
  exit 1
fi

export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}

if [[ -n "${TEST_DRYRUN:-}" ]]; then
  echo "[INFO] TEST_DRYRUN mode: used for testing"
  exit
fi

mkdir -p $EXP_DIR $LOG_DIR $CKPT_DIR
